from ultralytics import YOLO
import os
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description='训练YOLOv8水果识别模型')
    parser.add_argument('--data', type=str, default='data/data.yaml', help='数据配置文件路径')
    parser.add_argument('--epochs', type=int, default=100, help='训练轮数')
    parser.add_argument('--batch', type=int, default=16, help='批次大小')
    parser.add_argument('--imgsz', type=int, default=640, help='图像大小')
    parser.add_argument('--model', type=str, default='yolov8n.pt', help='模型类型')
    parser.add_argument('--device', type=str, default='', help='训练设备，如 0 或 0,1,2,3 或 cpu')
    return parser.parse_args()

def main():
    args = parse_args()
    
    # 确保模型目录存在
    os.makedirs('models', exist_ok=True)
    
    # 加载预训练模型
    model = YOLO(args.model)
    
    # 训练模型
    results = model.train(
        data=args.data,
        epochs=args.epochs,
        batch=args.batch,
        imgsz=args.imgsz,
        device=args.device,
        project='runs',
        name='train',
        exist_ok=True
    )
    
    # 验证模型
    model.val()
    
    # 复制最佳模型到models目录
    best_model_path = os.path.join('runs', 'train', 'weights', 'best.pt')
    if os.path.exists(best_model_path):
        import shutil
        shutil.copy(best_model_path, 'models/best.pt')
        print(f"最佳模型已复制到 models/best.pt")
    
    print("训练完成！")

if __name__ == '__main__':
    main() 